# -*- ecoding: utf-8 -*-
# @ModuleName: transform
# @Function: 
# @Author: Liweijian
# @Time: 2024/12/20 14:41

import numpy as np
import cv2
import os
import struct

def save_mnist_to_jpg(mnist_image_file, mnist_label_file, save_dir):
    if 'train' in os.path.basename(mnist_image_file):
        prefix = 'train'
    else:
        prefix = 'test'

    labelIndex = 0
    imageIndex = 0
    i = 0
    lbdata = open(mnist_label_file, 'rb').read()  # 二进制打开
    magic, nums = struct.unpack_from(">II", lbdata, labelIndex)  # 从 lbdata 中解析两个无符号整型：魔数和标签数量
    labelIndex += struct.calcsize('>II')  # 更新标签索引，指向下一个数据位置

    imgdata = open(mnist_image_file, "rb").read()
    magic, nums, numRows, numColumns = struct.unpack_from('>IIII', imgdata,imageIndex)  # 从 imgdata 中解析四个无符号整型：魔数、图像数量、行数和列数
    imageIndex += struct.calcsize('>IIII')  # 更新图像索引，指向下一个数据位置

    for i in range(nums):  # 遍历每一幅图像，循环次数为图像数量
        label = struct.unpack_from('>B', lbdata, labelIndex)[0]  # 从 lbdata 中读取一个无符号字节，表示当前图像的标签
        labelIndex += struct.calcsize('>B')  # 更新标签索引，指向下一个标签

        im = struct.unpack_from('>784B', imgdata, imageIndex)  # 从 imgdata 中读取 784 个无符号字节，表示 28x28 图像的像素值
        imageIndex += struct.calcsize('>784B')  # 更新图像索引，指向下一个图像数据的位置

        im = np.array(im, dtype='uint8')  # 将读取的像素值转换为 NumPy 数组，数据类型为无符号 8 位整数
        img = im.reshape(28, 28)  # 将一维数组重塑为 28x28 的二维数组，形成图像格式

        save_name = os.path.join(save_dir, '{}_{}_{}.jpg'.format(prefix, i, label))  # 生成保存图像的文件名，包括前缀、索引和标签
        cv2.imwrite(save_name, img)  # 使用 OpenCV 将图像保存为指定路径的 JPEG 文件


if __name__ == '__main__':
    train_images = 'D:/pythonProject/Final_MNIST/dataset/raw/train-images-idx3-ubyte'  # 训练集图像的文件名
    train_labels = 'D:/pythonProject/Final_MNIST/dataset/raw/train-labels-idx1-ubyte'  # 训练集label的文件名
    test_images = 'D:/pythonProject/Final_MNIST/dataset/raw/t10k-images-idx3-ubyte'  # 测试集图像的文件名
    test_labels = 'D:/pythonProject/Final_MNIST/dataset/raw/t10k-labels-idx1-ubyte'  # 测试集label的文件名

    save_train_dir = './data/train_images/'
    save_test_dir = './data/test_images/'

    if not os.path.exists(save_train_dir):
        os.makedirs(save_train_dir)
    if not os.path.exists(save_test_dir):
        os.makedirs(save_test_dir)

    save_mnist_to_jpg(test_images, test_labels, save_test_dir)
    save_mnist_to_jpg(train_images, train_labels, save_train_dir)